import sys
sys.path.append('C:/Users/Hu/Dropbox/Research/PythonWork/Cancer/src/STAT/')

import numpy as np
import csv
from datetime import datetime
from time import strftime
import matplotlib.pyplot as plt
import math
import random
from scipy import stats
import ols

'''
Global and Local Empirical Bayes Smoothers with Gamma Model
'''

def getCurTime():
    """
    get current time
    Return value of the date string format(%Y-%m-%d %H:%M:%S)
    """
    format='%Y-%m-%d %H:%M:%S'
    sdate = None
    cdate = datetime.now()
    try:
        sdate = cdate.strftime(format)
    except:
        raise ValueError
    return sdate

def build_data_list(inputCSV):
    sKey = []
    fn = inputCSV
    f = open(inputCSV)
    #ra = csv.DictReader(file(fn), dialect="excel")
    ra = csv.DictReader(f, dialect="excel")
    
    for record in ra:
        #print record[ra.fieldnames[0]], type(record[ra.fieldnames[-1]])
        for item in ra.fieldnames:
            temp = int(float(record[item]))
            sKey.append(temp)
    sKey = np.array(sKey)
    sKey.shape=(-1,len(ra.fieldnames))
    return sKey

def log(data):
    logdata = []
    for item in data:
        logdata.append(math.log(item))
    return logdata

def geneBpChoices(tmin, tmax, num):
    tmin = float(tmin)
    tmax = float(tmax)
    num = int(num)
    step = (tmax - tmin)/num
    tbpchoices = []
    for i in range(num):
        tbpchoices.append(tmin+i*step)
    return tbpchoices

def pickBreakpoint(response, x1, predictor):
    #print int(min(predictor))*10, int(max(predictor)+1)*10, int(max(predictor) - min(predictor) + 1)/2
    #bpChoices = geneBpChoices(min(predictor), max(predictor), 20)
    results = np.zeros((len(bpChoices)-1, 2))
    print bpChoices
    
    for i in range(len(bpChoices)-1):
        #print i
        x2star = (predictor - bpChoices[i]) * np.greater(predictor, bpChoices[i])   
        tempPredictor = np.array(zip(x1, predictor, x2star))
        #fileLoc = filePath + 'temp.csv'
        #np.savetxt(fileLoc, tempPredictor, delimiter=',', fmt = '%s')
        tempmodel = ols.ols(response, tempPredictor,'y',['F1F2', 'dist', 'diststar'])
        results[i,0] = i
        #results[i,1] = tempmodel.sse
        results[i,1] = tempmodel.R2

    optBP = int(results[np.argmax(results, axis = 0)[1],0])
    print 'Optimal Index:', optBP
    print 'Optimal changepoint: ', bpChoices[optBP], ' exp value: ', np.exp(bpChoices[optBP]), ' with R2 = ', results[optBP, 1]

    #x2star = (predictor - bpChoices[optBP]) * np.greater(predictor, bpChoices[optBP])
    #optPredictor = np.array(zip(predictor, x2star))
    #optmodel = ols.ols(response, optPredictor,'y',['x1', 'x2'])
    x2star = (predictor - bpChoices[optBP]) * np.greater(predictor, bpChoices[optBP])
    optPredictor = np.array(zip(x1, predictor, x2star))
    optmodel = ols.ols(response, optPredictor,'y',['F1F2', 'dist', 'diststar'])
    
    #return bpChoices[optBP], results, optmodel, optmodel.b[0]+optmodel.b[1]*predictor+optmodel.b[2]*x2star
    print results, optmodel.b
    print optmodel.summary()
    return results

def pickBreakpointV2(response, x1, predictor):
    #print int(min(predictor))*10, int(max(predictor)+1)*10, int(max(predictor) - min(predictor) + 1)/2
    #bpChoices = geneBpChoices(min(predictor), max(predictor), 20)
    results = np.zeros((len(bpChoices)-1, 2))
    print bpChoices
    
    for i in range(len(bpChoices)-1):
        #print i
        x2star = (predictor - bpChoices[i]) * np.greater(predictor, bpChoices[i])
        x1star = x1 * np.greater(predictor, bpChoices[i]) 
        tempPredictor = np.array(zip(x1, x1star, predictor, x2star))
        #fileLoc = filePath + 'temp.csv'
        #np.savetxt(fileLoc, tempPredictor, delimiter=',', fmt = '%s')
        tempmodel = ols.ols(response, tempPredictor,'y',['F1F2', 'F1F2star', 'dist', 'diststar'])
        results[i,0] = i
        #results[i,1] = tempmodel.sse
        results[i,1] = tempmodel.R2

    optBP = int(results[np.argmax(results, axis = 0)[1],0])
    print 'Optimal Index:', optBP
    print 'Optimal changepoint: ', bpChoices[optBP], ' exp value: ', np.exp(bpChoices[optBP]), ' with R2 = ', results[optBP, 1]

    #x2star = (predictor - bpChoices[optBP]) * np.greater(predictor, bpChoices[optBP])
    #optPredictor = np.array(zip(predictor, x2star))
    #optmodel = ols.ols(response, optPredictor,'y',['x1', 'x2'])
    x1star = x1 * np.greater(predictor, bpChoices[i])
    x2star = (predictor - bpChoices[optBP]) * np.greater(predictor, bpChoices[optBP])
    optPredictor = np.array(zip(x1, x1star, predictor, x2star))
    optmodel = ols.ols(response, optPredictor,'y',['F1F2', 'F1F2star', 'dist', 'diststar'])
    
    #return bpChoices[optBP], results, optmodel, optmodel.b[0]+optmodel.b[1]*predictor+optmodel.b[2]*x2star
    print results, optmodel.b
    print optmodel.summary()
    return results

#--------------------------------------------------------------------------
#MAIN

if __name__ == "__main__":
    print '===================================================='
    print "begin at " + getCurTime()
    #print geneBpChoices(1.4, 7.9, 20)

    filePath = 'C:/_DATA/migration_census_2000/'
    file = filePath + 'census_county_migration_format_dist.csv' #[dist, orgid, desid, flowvol]
    data = build_data_list(file)

    inoutflowfile = filePath + 'census_county_inoutflow.csv'    #[inflow, outflow]
    inoutflow = build_data_list(inoutflowfile)

    x1 = []
    for item in data:
        x1.append(np.log(float(inoutflow[item[1],1])*inoutflow[item[2],0]))

    #bpChoices = geneBpChoices(min(data[:,0]), max(data[:,0]), 20)
    #bpChoices = geneBpChoices(100000, 3000000, 100)
    bpChoices = range(100000, 500000, 5000)
    print bpChoices
    bpChoices = np.log(bpChoices)

    result = pickBreakpointV2(np.log(data[:,-1]), x1, np.log(data[:,0]))

    #fileLoc = filePath + 'segmented_data.csv'
    #np.savetxt(fileLoc, zip(np.log(ajustOB50), np.log(data[:,0])), delimiter=',', fmt = '%s')
    